package com.lingtianyu.tech.springcloud.grayzone.grey.loadbalancer;

import com.lingtianyu.tech.springcloud.grayzone.grey.context.GreyContext;
import io.vavr.Tuple;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.discovery.DiscoveryClient;
import org.springframework.cloud.client.loadbalancer.*;
import org.springframework.cloud.client.loadbalancer.reactive.ReactiveLoadBalancer;
import org.springframework.cloud.loadbalancer.blocking.client.BlockingLoadBalancerClient;
import org.springframework.cloud.loadbalancer.support.LoadBalancerClientFactory;

import java.io.IOException;
import java.util.*;

import static org.springframework.cloud.client.loadbalancer.reactive.ReactiveLoadBalancer.REQUEST;

/**
 * 此处在需要的时候注入，则可以开启新的灰度负载均衡能力
 */
@Slf4j
public class GreyLoadBalancerClient extends BlockingLoadBalancerClient {
    public static final String ZONE_NAME = "lingtianyu-zone";
    public static final String DEFAULT_ZONE = "default";
    private static Random RANDOM = new Random();
    private final ReactiveLoadBalancer.Factory<ServiceInstance> loadBalancerClientFactory;
    private final DiscoveryClient discoveryClient;

    public GreyLoadBalancerClient(LoadBalancerClientFactory loadBalancerClientFactory, LoadBalancerProperties properties, DiscoveryClient discoveryClient) {
        super(loadBalancerClientFactory, properties);
        this.loadBalancerClientFactory = loadBalancerClientFactory;
        this.discoveryClient = discoveryClient;
    }

    public GreyLoadBalancerClient(ReactiveLoadBalancer.Factory<ServiceInstance> loadBalancerClientFactory, DiscoveryClient discoveryClient) {
        super(loadBalancerClientFactory);
        this.loadBalancerClientFactory = loadBalancerClientFactory;
        this.discoveryClient = discoveryClient;
    }

    public <T> T execute(String serviceId, LoadBalancerRequest<T> request) throws IOException {
        String hint = getHint(serviceId);
        LoadBalancerRequestAdapter<T, DefaultRequestContext> lbRequest = new LoadBalancerRequestAdapter<>(request,
                new DefaultRequestContext(request, hint));
        Set<LoadBalancerLifecycle> supportedLifecycleProcessors = getSupportedLifecycleProcessors(serviceId);
        supportedLifecycleProcessors.forEach(lifecycle -> lifecycle.onStart(lbRequest));
        ServiceInstance serviceInstance = this.choose(serviceId, lbRequest);
        if (serviceInstance == null) {
            supportedLifecycleProcessors.forEach(lifecycle -> lifecycle.onComplete(
                    new CompletionContext<>(CompletionContext.Status.DISCARD, lbRequest, new EmptyResponse())));
            throw new IllegalStateException("No instances available for " + serviceId);
        }
        return execute(serviceId, serviceInstance, lbRequest);
    }

    public <T> T execute(String serviceId, ServiceInstance serviceInstance, LoadBalancerRequest<T> request) throws IOException {
        return super.execute(serviceId, serviceInstance, request);
    }

    @Override
    public ServiceInstance choose(String serviceId) {
        return choose(serviceId, REQUEST);
    }

    @Override
    public <T> ServiceInstance choose(String serviceId, Request<T> request) {
        List<ServiceInstance> list = discoveryClient.getInstances(serviceId);
        if (list.isEmpty()) {
            if (log.isWarnEnabled()) {
                log.warn("No servers available for service: " + serviceId);
            }
            return null;
        }

        String springCloudZone = GreyContext.getSpringCloudZone();

        Map<String, List<ServiceInstance>> zoneServiceMap = io.vavr.collection.List.ofAll(list).groupBy(serviceInstance -> {
            Map<String, String> metadata = serviceInstance.getMetadata();
            if (metadata == null || metadata.isEmpty()) {
                return DEFAULT_ZONE;
            }
            String springcloudZone = metadata.get(ZONE_NAME);
            if (springcloudZone == null || springcloudZone.isEmpty()) {
                return DEFAULT_ZONE;
            }
            return springcloudZone;
        }).map((k, v) -> Tuple.of(k, v.toJavaList())).toJavaMap();
        //如果服务只有default节点，则走原本的路由逻辑
        if (zoneServiceMap.size() == 1 && zoneServiceMap.get(DEFAULT_ZONE) != null) {
//            List<ServiceInstance> defaultService = zoneServiceMap.getOrDefault(DEFAULT_ZONE, Collections.emptyList());
            return super.choose(serviceId, request);
        }
        //否则按照灰度规则取
        List<ServiceInstance> serviceInstances = zoneServiceMap.get(springCloudZone);
        if (serviceInstances == null || serviceInstances.isEmpty()) {
            List<ServiceInstance> defaultService = zoneServiceMap.getOrDefault(DEFAULT_ZONE, Collections.emptyList());
            return getHostByRandomWeight(defaultService);
        }
        return getHostByRandomWeight(serviceInstances);
    }

    private Set<LoadBalancerLifecycle> getSupportedLifecycleProcessors(String serviceId) {
        return LoadBalancerLifecycleValidator.getSupportedLifecycleProcessors(
                loadBalancerClientFactory.getInstances(serviceId, LoadBalancerLifecycle.class),
                DefaultRequestContext.class, Object.class, ServiceInstance.class);
    }

    private String getHint(String serviceId) {
        LoadBalancerProperties properties = loadBalancerClientFactory.getProperties(serviceId);
        String defaultHint = properties.getHint().getOrDefault("default", "default");
        String hintPropertyValue = properties.getHint().get(serviceId);
        return hintPropertyValue != null ? hintPropertyValue : defaultHint;
    }

    public static ServiceInstance getHostByRandomWeight(List<ServiceInstance> serviceInstances) {
        if(serviceInstances == null || serviceInstances.isEmpty()) {
            return null;
        }
        int i = RANDOM.nextInt(serviceInstances.size());
        return serviceInstances.get(i);
    }
}
